from calendar import c
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt

def get_bn_diff(path):
    weights = torch.load(path, weights_only=False)["state_dict"]

    student_running_mean = []
    student_running_var = []
    teacher_running_mean = []
    teacher_running_var = []

    for k, v in weights.items():
        if "sed_student" in k:
            if "running_mean" in k:
                student_running_mean.append(v)
            elif "running_var" in k:
                student_running_var.append(v)
        elif "sed_teacher" in k:
            if "running_mean" in k:
                teacher_running_mean.append(v)
            elif "running_var" in k:
                teacher_running_var.append(v)

    mean_diff = []
    var_diff = []
    for layer in range(len(student_running_mean)):
        mean_diff.append(F.mse_loss(student_running_mean[layer], teacher_running_mean[layer]).item())
        var_diff.append(F.mse_loss(student_running_var[layer], teacher_running_var[layer]).item())
    return mean_diff, var_diff

mean_1, var_1 = get_bn_diff("/home/shaonian/SED/sssl_sed/codes/exp/da_exp_loss/baseline/Aug_view1+view2/BN_fixed/CL_Stu-Stu/version_0/ckpts/epoch=289-obj_metric=1.454.ckpt")
mean_2, var_2 = get_bn_diff("/home/shaonian/SED/sssl_sed/codes/exp/da_exp_loss/baseline_reproduce/version_0/ckpts/epoch=289-obj_metric=1.457.ckpt")

plt.plot(mean_1, label="mean_1", color="red")
plt.plot(mean_2, label="mean_bs", color="blue")
plt.legend()
plt.xlabel("layers")
plt.ylabel("MSE")
plt.savefig("bn_mean_diff.png")
plt.close()

plt.plot(var_1, "--*", label="var_1", color="red")
plt.plot(var_2, "--*", label="var_bs", color="blue")
plt.legend()
plt.xlabel("layers")
plt.ylabel("MSE")

plt.title("BN diff")
plt.savefig("bn_diff.png")